#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RKNN模型信息检查工具 - 仅结构分析
不执行推理，只分析模型的输入输出结构信息
"""

import os
import struct

def parse_rknn_model_info(model_path):
    """解析RKNN模型文件的基本信息"""
    if not os.path.exists(model_path):
        print(f"❌ 模型文件不存在: {model_path}")
        return False
    
    file_size = os.path.getsize(model_path)
    print(f"📁 模型文件: {model_path}")
    print(f"📏 文件大小: {file_size / (1024*1024):.2f} MB")
    
    try:
        with open(model_path, 'rb') as f:
            # 读取文件头
            header = f.read(16)
            if len(header) < 16:
                print("❌ 文件过小，不是有效的RKNN模型")
                return False
            
            # 简单的文件格式检查
            if b'RKNN' in header[:16]:
                print("✅ 检测到RKNN文件格式标识")
            else:
                print("⚠️  未找到RKNN格式标识，但可能仍是有效模型")
            
            return True
            
    except Exception as e:
        print(f"❌ 读取模型文件失败: {e}")
        return False

def estimate_model_specs(model_path):
    """根据文件大小估算模型规格"""
    file_size = os.path.getsize(model_path)
    
    print(f"\n📊 模型规格估算（基于文件大小 {file_size/(1024*1024):.1f}MB）:")
    
    if file_size < 10 * 1024 * 1024:  # < 10MB
        print("🔹 轻量级模型 (< 10MB)")
        print("   - 可能是YOLOv8n或类似轻量版本")
        print("   - 推理速度快，精度相对较低")
    elif file_size < 50 * 1024 * 1024:  # < 50MB  
        print("🔹 中等规模模型 (10-50MB)")
        print("   - 可能是YOLOv8s/m版本")
        print("   - 平衡的速度和精度")
    elif file_size < 200 * 1024 * 1024:  # < 200MB
        print("🔹 大型模型 (50-200MB)")
        print("   - 可能是YOLOv8l版本")
        print("   - 高精度，推理速度相对较慢")
    else:
        print("🔹 超大型模型 (> 200MB)")
        print("   - 可能是YOLOv8x或自定义大模型")
        print("   - 最高精度，需要更多计算资源")

def analyze_rknn_without_inference(model_path):
    """不执行推理的RKNN模型分析"""
    print("=" * 50)
    print("🔍 RKNN模型结构分析（无推理）")
    print("=" * 50)
    
    # 基本文件信息
    if not parse_rknn_model_info(model_path):
        return False
    
    # 估算模型规格
    estimate_model_specs(model_path)
    
    # 基于你的训练信息提供分析
    print(f"\n📋 基于训练配置的预期信息:")
    print(f"   - 输入尺寸: 根据转换脚本应为 [1, 3, 640, 640] 或 [1, 3, 480, 640]")
    print(f"   - 输出格式: [1, 19, 6300] (4坐标 + 15类别)")
    print(f"   - 数据类型: FP16")
    print(f"   - 类别数量: 15种物资")
    
    print(f"\n💡 建议的验证步骤:")
    print(f"1. 在RK3588硬件上加载模型验证实际输入输出格式")
    print(f"2. 使用C++代码测试推理性能") 
    print(f"3. 调整置信度阈值以获得最佳检测效果")
    
    return True

def main():
    """主函数"""
    print("🔍 RKNN模型信息检查工具")
    print("📌 此工具仅分析模型文件结构，不执行推理")
    print("=" * 50)
    
    # 查找RKNN模型文件
    rknn_files = []
    search_paths = [
        'wuzi.rknn',
        'wuzi_640x640.rknn', 
        'wuzi_fixed.rknn',
        'model/wuzi.rknn',
        'model/wuzi_640x640.rknn',
        'model/wuzi_fixed.rknn'
    ]
    
    for path in search_paths:
        if os.path.exists(path):
            rknn_files.append(path)
    
    if not rknn_files:
        print("❌ 未找到RKNN模型文件")
        print("查找路径:")
        for path in search_paths:
            print(f"  - {path}")
        print("\n建议:")
        print("1. 确保已运行模型转换脚本")
        print("2. 检查模型文件路径")
        return
    
    print(f"✅ 找到RKNN模型: {rknn_files}")
    
    # 分析每个模型
    for model_file in rknn_files:
        print(f"\n🔍 分析模型: {model_file}")
        success = analyze_rknn_without_inference(model_file)
        
        if success:
            print(f"✅ {model_file} 信息分析完成")
        else:
            print(f"❌ {model_file} 分析失败")
        
        print("-" * 50)
    
    print("\n🎯 下一步建议:")
    print("1. 使用PC端ONNX验证:")
    print("   python test_onnx_model.py")
    print("2. 在RK3588上编译并测试:")
    print("   ./build.sh && ./build_and_test.sh")
    print("3. 如果检测效果不佳:")
    print("   - 调整 include/det/postprocess.h 中的阈值")
    print("   - 检查训练数据与实际场景的匹配度")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n⚠️  用户中断")
    except Exception as e:
        print(f"\n❌ 程序异常: {e}") 